from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
import os
import numpy as np

from others.utils import parse_voc_annotation

voc_data_dir = '../VOCdevkit'
year = '2007'

# 创建 COCO 格式的数据结构
coco = COCO()

# 分别处理图像和标注
# 读取图像信息
image_dir = os.path.join(voc_data_dir, 'VOC' + year, 'JPEGImages')
for img_name in os.listdir(image_dir):
    if img_name.endswith('.jpg'):
        img_path = os.path.join(image_dir, img_name)
        img_id = int(os.path.splitext(img_name)[0])
        width, height = 512, 512  # 假设的图像尺寸
        coco.addImg(img_id, img_path, width=width, height=height)

# 读取标注信息
annotations_dir = os.path.join(voc_data_dir, 'VOC' + year, 'Annotations')
for ann_name in os.listdir(annotations_dir):
    if ann_name.endswith('.xml'):
        ann_path = os.path.join(annotations_dir, ann_name)
        ann_id = int(os.path.splitext(ann_name)[0])
        bbox, category_id = parse_voc_annotation(ann_path)  # 假设有函数可以解析 VOC 标注格式
        area = (bbox[2] - bbox[0]) * (bbox[3] - bbox[1])
        iscrowd = 0
        annotation = {'image_id': img_id, 'bbox': bbox, 'category_id': category_id, 'area': area, 'iscrowd': iscrowd}
        coco.addAnn(annotation)

# 创建类别的映射
categories = ['person', 'dog', 'cat']  # 假设你已经知道了类别
for i, cat in enumerate(categories):
    coco.addCat(i, cat)

# 保存 COCO 格式的数据
coco.saveAnns()
coco.saveCats()

# 验证 COCO 格式数据的正确性
coco.dataset['images'][0]
coco.dataset['categories'][0]
coco.dataset['annotations'][0]